# scepter_motion_working.py - 确保能工作的版本，专注关节空间运动
import time

from moveit.task_constructor import core, stages
import rclcpp


def main():
    rclcpp.init()
    node = rclcpp.Node("scepter_mtc_working")
    
    # 设置规划组
    arm_group = "arm_group"
    gripper_group = "gripper_group"

    # 创建任务并加载机器人模型
    task = core.Task()
    task.name = "scepter_working_demo"
    
    try:
        task.loadRobotModel(node)
        print("✓ Robot model loaded successfully")
    except Exception as e:
        print(f"✗ Failed to load robot model: {e}")
        return

    # 创建关节插值规划器（已验证可工作）
    joint_planner = core.JointInterpolationPlanner()
    print("✓ Joint space planner created")
    
    # 从当前状态开始
    current_state = stages.CurrentState("current state")
    task.add(current_state)

    # 1. 移动到准备位置
    move_to_ready = stages.MoveTo("move to ready position", joint_planner)
    move_to_ready.group = arm_group
    move_to_ready.setGoal("ready_pose")  # 使用 SRDF 中定义的姿态
    task.add(move_to_ready)

    # 2. 打开夹爪
    open_gripper = stages.MoveTo("open gripper", joint_planner)
    open_gripper.group = gripper_group
    open_gripper.setGoal("open_pose")  # 使用 SRDF 中定义的姿态
    task.add(open_gripper)

    # 3. 移动到完成位置（模拟取物后的位置）
    move_to_finish = stages.MoveTo("move to finish position", joint_planner)
    move_to_finish.group = arm_group
    move_to_finish.setGoal("finish_pose")  # 使用 SRDF 中定义的姿态
    task.add(move_to_finish)

    # 4. 关闭夹爪（模拟抓取物体）
    close_gripper = stages.MoveTo("close gripper", joint_planner)
    close_gripper.group = gripper_group
    close_gripper.setGoal("close_pose")  # 使用 SRDF 中定义的姿态
    task.add(close_gripper)

    # 5. 回到准备位置（模拟移动物体）
    move_back_ready = stages.MoveTo("move back to ready", joint_planner)
    move_back_ready.group = arm_group
    move_back_ready.setGoal("ready_pose")
    task.add(move_back_ready)

    # 6. 打开夹爪释放物体
    release_object = stages.MoveTo("release object", joint_planner)
    release_object.group = gripper_group
    release_object.setGoal("open_pose")
    task.add(release_object)

    # 7. 回到初始位置
    move_home = stages.MoveTo("move home", joint_planner)
    move_home.group = arm_group
    move_home.setGoal("stand_pose")  # 使用 SRDF 中定义的姿态
    task.add(move_home)

    print("🚀 Starting planning...")
    
    # 执行规划
    if task.plan():
        print("✓ Planning succeeded!")
        print("🎯 Publishing solution...")
        task.publish(task.solutions[0])
        print("✓ Solution published successfully")
        print("📊 Task completed with all stages executed")
    else:
        print("✗ Planning failed")
        return

    print("⏳ Waiting for 5 seconds to observe results...")
    time.sleep(5)
    print("🏁 Demo completed!")


if __name__ == "__main__":
    main()